#include <iostream>
#include <vector>

#include "../include/tensor.h"
using namespace std;

int main()
{
    // 随机初始化 Tensor
    Tensor random_tensor({2, 3, 4}, -1.0, 1.0);
    cout << random_tensor << endl;
    cout << "==========================================" << endl;

    // 从 vector 初始化 Tensor
    std::vector<double> data = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0};
    Tensor vector_tensor(data, {2, 3, 2});
    cout << vector_tensor << endl;
    cout << "==========================================" << endl;

    // 多维索引访问
    std::cout << "Element at indices [1, 2, 1]: " << vector_tensor.at({1, 2, 1}) << std::endl;
    std::cout << "Element at indices [1, 1, 1]: " << vector_tensor[{1, 1, 1}] << std::endl;
    cout << "==========================================" << endl;

    // 改变视图
    Tensor view_tensor = vector_tensor.view({6, 2});
    cout << view_tensor << endl;
    cout << "==========================================" << endl;

    // +
    Tensor c = Tensor(vector<double>({2.0})) + Tensor(vector<double>({3.0}));
    cout << c << endl;
    cout << "==========================================" << endl;

    // *
    Tensor d = Tensor(vector<double>({2.0})) * Tensor(vector<double>({3.0})) + Tensor(vector<double>({4.0})) * Tensor(vector<double>({5.0}));
    d.backward(); // ok
    auto prev = d.getData()[0]->_prev;
    for (auto v : prev)
    {
        cout << v->data << "\t grad:" << v->grad << endl;
        auto p = v->_prev;
        for (auto vv : p)
        {
            cout << vv->data << "\t grad:" << vv->grad << endl;
        }
    }

    // sum
    Tensor e = Tensor({1, 2, 3, 4, 5, 6}, {2, 3});
    Tensor f = e.sum(0);
    f.backward();
    cout << e.getData()[0]->grad << endl;
    cout << "==========================================" << endl;
}